/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2023. All rights reserved.
 *
 * Function : z = x + y
 * This sample is a very basic sample that implements vector add on Ascend plaform.
 */
#include "kernel_operator.h"
#include "math.h"
#include <float.h>
#include<stdio.h>
// tensor num for each queue
using namespace AscendC;

template<typename TYPE_X, typename TYPE_Y> class KernelAdd {
    using T = TYPE_X;
public:
    __aicore__ inline KernelAdd() {}
    __aicore__ inline void Init(GM_ADDR var, GM_ADDR indices,GM_ADDR updates,
                                int32_t* input_shape, int32_t* update_shape, 
                                int32_t totalLengthInput, int32_t totalLengthUpdate,
                                int32_t dim, int32_t mode, int32_t dimNum)
    {
        this->input_shape = input_shape;
        this->update_shape = update_shape;
        this->totalLengthInput = totalLengthInput;
        this->totalLengthUpdate = totalLengthUpdate;
        this->dim = dim;
        this->mode = mode;
        this->dimNum = dimNum;
        
        
        varGm.SetGlobalBuffer((__gm__ TYPE_X *)var, this->totalLengthInput);
        indicesGm.SetGlobalBuffer((__gm__ int32_t *)indices, this->totalLengthUpdate);
        updatesGm.SetGlobalBuffer((__gm__ TYPE_X *)updates, this->totalLengthUpdate);

        // pipe.InitBuffer(tmp1, sizeof(float));
        // pipe.InitBuffer(tmp2, sizeof(float));
        // pipe.InitBuffer(sum_buffer, 10240*sizeof(float));
        // pipe.InitBuffer(work_buffer, 10240*sizeof(float));
    }

    // __aicore__  int32_t input_to_output(int32_t input_index)
    // {
    //     return 0;
    // }


    __aicore__ inline void print_arr(int32_t* array, int32_t arr_len){
        for(int32_t i=0; i<arr_len; ++i)
        {
            printf("%d ", array[i]);
            
        }
        printf("\n");
    }

    __aicore__ inline void Process()
    {
        for(int index = 0;index<this->totalLengthUpdate;++index)
        {
            int32_t update_indexes[100];
            int32_t index_to_reduce = index;
            for(int i = 0;i<this->dimNum;++i)
            {
                // printf("%d ", this->update_shape[dimNum-i-1]);
                update_indexes[i] = index_to_reduce%this->update_shape[dimNum-i-1];
                index_to_reduce = index_to_reduce/this->update_shape[dimNum-i-1];
            }

            // print_arr(update_indexes, this->dimNum);


            int32_t indice = indicesGm.GetValue(index);
            TYPE_X update = updatesGm.GetValue(index);
            update_indexes[dimNum-this->dim-1] = indice;

            int32_t output_index = 0;
            for(int i = 0;i<this->dimNum;++i)
            {
                output_index = output_index*this->input_shape[i];
                output_index+=update_indexes[dimNum-i-1];
                
            }

            if (mode==0)
                varGm.SetValue(output_index, update);
            else{
                float init_value = float(varGm.GetValue(output_index));
                float final_value;
                if (mode==1)
                {
                    final_value = init_value+float(update);
                }
                else
                {
                    final_value = init_value*float(update);
                }

                varGm.SetValue(output_index, (TYPE_X)final_value);
            }
        }
        
    }



    __aicore__ inline void Process_vec()
    {

    }
private:
    TPipe pipe;
    GlobalTensor<TYPE_X> varGm;
    GlobalTensor<int32_t> indicesGm;
    GlobalTensor<TYPE_Y> updatesGm;


    // AscendC::TBuf<AscendC::QuePosition::VECCALC> tmp1, tmp2;
    // AscendC::TBuf<AscendC::QuePosition::VECCALC> sum_buffer, work_buffer;
    int32_t* input_shape, *update_shape;
    int32_t totalLengthInput, totalLengthUpdate, dim, mode, dimNum;
};

extern "C" __global__ __aicore__ void scatter_elements(GM_ADDR var, GM_ADDR indices,GM_ADDR updates, GM_ADDR workspace, GM_ADDR tiling)
{

    GET_TILING_DATA(tiling_data, tiling);
    
    if constexpr (std::is_same_v<DTYPE_VAR, u_int8_t>)
    {
        KernelAdd<int8_t, int8_t> op;
        op.Init(var, indices, updates, tiling_data.input_shape, tiling_data.update_shape, tiling_data.totalLengthInput, tiling_data.totalLengthUpdate, tiling_data.dim, tiling_data.mode, tiling_data.dimNum);  
        op.Process();
    }
    else
    {

        KernelAdd<DTYPE_VAR, DTYPE_VAR> op;
        op.Init(var, indices, updates, tiling_data.input_shape, tiling_data.update_shape, tiling_data.totalLengthInput, tiling_data.totalLengthUpdate, tiling_data.dim, tiling_data.mode, tiling_data.dimNum);  
        // if (tiling_data.mode==0)
        //     op.Process_vec();
        // else
        op.Process();
    }

}
